from math import sin, acos
import os
import pysam
import subprocess
from allfunctions import dimer, revcomp
from random import shuffle

# for repeat counting
xmer= 12
# what are the four bases
bases= ['a','c','g','t']

class memoized(object):
    """Decorator that caches a function's return value each time it is called.
    If called later with the same arguments, the cached value is returned, and
    not re-evaluated.
    """
    def __init__(self, func):
       self.func = func
       self.cache = {}
    def __call__(self, *args):
       try:
          return self.cache[args]
       except KeyError:
          value = self.func(*args)
          self.cache[args] = value
          return value
       except TypeError:
          # uncachable -- for instance, passing a list as an argument.
          # Better to not cache than to blow up entirely.
          return self.func(*args)
    def __repr__(self):
       """Return the function's docstring."""
       return self.func.__doc__
    def __get__(self, obj, objtype):
       """Support instance methods."""
       return functools.partial(self.__call__, obj)


@memoized
def repeat_counter(gpath, s, xmer):
    # the codec
    code= {'A': 0, 'C': 1, 'G': 2, 'T': 3, 'a': 0, 'c': 1, 'g': 2, 't': 3}
    # to be used with the repeatDB
    # the returned value is the number of s in the genome
    input= open(os.path.join(gpath, '%smer_repeats.txt' %xmer))
    title = input.readline()
    xmer, lenmax= title.strip().split('\t')
    xmer = int(xmer)
    lenmax = int(lenmax)

    score= 0
    s = list(s)
    read= input.read
    strip= str.strip
    for i in xrange(len(s[:-xmer+1])):
        # all but the last xmer-1 nucleotides
        temp= s[i:i+xmer]
        temp.reverse()
        R= [code[p]*4**j for p,j in zip(temp, xrange(xmer))]
        # each line has a "\t" and "\n"
        value= len(title) + sum(R) * (xmer+lenmax+2) + (xmer+1)
        # take the head to the position you want
        input.seek(value)
        # read from that position only the count
        data= read(lenmax)
        # add on top 
        score+= int(strip(data))
    return score

                
@memoized
def score_repeats(gpath, s, xmer):
    # max value should be 1
    # so scale everything accordingly (0-1)
    # length adjusted repeat count should be<2000 (Li Genome Res, 2009)
    # for each xmer if the count is 1, that is as good as it gets
    score= repeat_counter(gpath, s, xmer) / len(s)
    # anything larger that 2000 should be maxed-out at 1
    adj_score= min(1, score/2000.0)
    # what is the probability of finding an xmer in the genome
    # value= (2.0 * 3000000000 / 4**xmer)
    return sin(acos(adj_score))


def query_tabix(fpath, Chr, Start, Stop):
    try:
        return [p.split("\t") for p in fpath.fetch(Chr, Start, Stop)]
    except AttributeError:
        return list()


def find_exon(exon_file, Chr, Start, Stop):
    # get tabix output
    query = query_tabix(exon_file, Chr, Start, Stop)

    exons_list = set()

    if len(query):
        for line in query:
            junk, Acc_number, Chr, strand, Gstart, Gstop, Cstart, Cstop, Exoncount, Exon_start, Exon_stop, junk, Gene_name, status, junk, Frame = line

            # proper formating
            Exoncount= int(Exoncount)
            Exon_start= Exon_start.strip().split(',')
            Exon_stop= Exon_stop.strip().split(',')
            Frame= Frame.strip().split(',')
            del Exon_start[-1]
            del Exon_stop[-1]
            del Frame[-1]

            # strand work
            # fix the start coordinate of the plus strand ORF
            Cstart= str(int(Cstart)+1)

            if status=='cmpl':
                if strand == '+':
                    for i in xrange(Exoncount):
                        Estart = int(Exon_start[i]) + 1
                        Estop = int(Exon_stop[i])

                        if Estart <=  Start <=  Stop <= Estop or Start <= Estart <= Stop <= Estop or Estart <= Start <= Estop <= Stop or Start <= Estart < Estop <= Stop:
                            exons_list.add("%s:%s" %(Gene_name, i+1))
                            break
                        elif Stop < Estart:
                            break

                else:
                    Exon_start.reverse()
                    Exon_stop.reverse()
                    Frame.reverse()

                    for i in xrange(Exoncount):
                        Estart = int(Exon_start[i]) + 1
                        Estop = int(Exon_stop[i])

                        if Estart <=  Start <=  Stop <= Estop or Start <= Estart <= Stop <= Estop or Estart <= Start <= Estop <= Stop or Start <= Estart < Estop <= Stop:
                            exons_list.add("%s:%s" %(Gene_name, i+1))
                            break
                        elif Start > Estop:
                            break

            else:
                exons_list.add(Gene_name)
            
    exons_list = sorted(exons_list)
    
    if len(exons_list):
        return '|'.join(exons_list)
    else:
        return "unk"


def check_snp(snp_file, Chr, Start, Stop, strand):
    # get tabix output
    query = query_tabix(snp_file, Chr, Start, Stop)

    if strand == "F":
        midStart = Start + int(((Stop - Start) + 1) * 1.0/3)
        for line in query:
            snpS = int(line[2])
            snpE = int(line[3])
            if snpE != snpS:
                # if deletion than the values are the same
                snpS += 1
            if Start<=snpS<=snpE<=Stop or snpS<=Start<=snpE<=Stop:
                # position of the overlap is critical,
                # allow 5' mismatches
                if snpE > midStart:
                    return "FAIL"
            elif Start <= snpS <= Stop <= snpE:
                return "FAIL"            
        return 'PASS'       

    else:
        midStart = Start + int(((Stop - Start) + 1) * 2.0/3)
        for line in query:
            snpS = int(line[2])
            snpE = int(line[3])
            if snpE != snpS:
                # if deletion than the values are the same
                snpS += 1

            if Start <= snpS <= snpE <= Stop or Start <= snpS <= Stop <= snpE:
                # position of the overlap is critical, allow 5' mismatches
                if snpS > midStart:
                    return "FAIL"
            elif snpS <= Start <= snpE <= Stop:
                return "FAIL"

        return 'PASS'


def check_repeats(F, R, path):
    # ref: Li Genome Res, 2009
    score = 0
    for each in (F, R):
        score+= score_repeats(path, each, xmer)
    return score


def olcount(s, pattern):
    """Returns how many pattern on s, works for overlapping"""
    count= 0
    x= 0
    while 1:
        try:
            i= s.index(pattern,x)
        except ValueError:
            break
        count+= 1
        x= i+1
    return count


def Penalty(F, R): #penalty score for various factors
    # 3' GC content 
    cutoff= 3 # no more than 3
    penalty= 0.0
    if F[-5:].count('g') + F[-5:].count('c') > cutoff:
        penalty-= 0.5
    if R[-5:].count('g') + R[-5:].count('c') > cutoff:
        penalty-= 0.5
    # dinucleotide stretch
    cutoff= 4 # no more than 3
    base_each= 1
    while base_each <= len(bases):
        bases_each= 1
        while bases_each <= (len(bases)-1):
            dinuc= (bases[0] + bases[bases_each]) * cutoff
            if dinuc in F:
                penalty-= (0.5 * F.count(dinuc))
            if dinuc in R:
                penalty-= (0.5 * R.count(dinuc))
            bases_each+= 1
        bases.insert(4, bases[0])
        del bases[0]
        base_each+= 1
    # single nucleotide runs
    cutoff= 4 # no more than 3
    base_each= 0
    while base_each < len(bases):
        base_run= bases[base_each] * cutoff
        if olcount(F, base_run) > 0:
            penalty-= (0.5 * olcount(F, base_run))
        if olcount(R,base_run)>0:
            penalty-= (0.5 * olcount(R, base_run))
        base_each+= 1
    return penalty
     

def LenScore(F, R):
    if F == R == 20:
        return 0.5
    return 0


# this is the picker function 
# the function input is a list
def pcr_picker(s, REdict, PriF, Middle, PriR, Arm_max, pntcutoff, pick_reverse, TTemp2, even, path, Picked_code, gfserver, Remove=False):

    s.sort(reverse=True)
    # now sort desending
    # here implement dimer check for the pair
    # each F and R primer is dimer checked during design,
    # which is the fastest since number of primer is reduced upfront

    len_TTemp2= len(TTemp2)
    if even:
        strand= 'Tube1'
    else:
        strand= 'Tube2'
                    
    for pair in s:
        goodpair= True
        # you cannot use any of the F/R primers you picked before
        if pair[11] not in Picked_code and pair[1] not in Picked_code:
            if dimer(pair[2],pair[12],pntcutoff=pntcutoff)=='PASS':
                Picked_code[pair[1]]= 0
                Picked_code[pair[11]]= 0
                thepair= pair[:] + [strand]

                if not gfserver == "false":           
                    # in-silico PCR
                    # this does not work on my mac only on linux
                    # you can define -maxSize=max_len*5 (default=4000)
                    # which should be ok
                    proc = subprocess.Popen("%s %s %s stdout -out=bed" %(gfserver, pair[2], pair[12]), shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                    matches = [p for p in proc.stdout]
                    if len(matches) != 1:
                        # there has to be one and only one match
                        goodpair = False
                        #print matches
                    
                elif Remove:
                    goodpair= False
                    repeat_score= check_repeats(pair[2],pair[12], path)
                    if repeat_score > 0:
                        pair[0]+= repeat_score
                        goodpair= True
                        
                if goodpair:
                    if pick_reverse == None:
                        keep= True
                        for olds in TTemp2[:-1]:
                            if olds[9] < thepair[8] < thepair[9] < olds[18] < olds[19] < thepair[18] or olds[19] < thepair[8] or thepair[9] < olds[8] < olds[19] < thepair[18]:
                                keep= True
                            else:
                                keep= False
                                break
                        if keep:
                            return Picked_code, thepair
                    else:
                        return Picked_code, thepair
    return Picked_code, 'FAIL'


def pcr_TmScore(F, R):#max score 2
    Score= 0
    if 60 >= F >= 55 and 60 >= R >= 55:
        Score+= 1.0
    if abs(F-R) <= 5:
        Score+= 1.0
    return Score


def pcr_GCScore(F, R): #max score 2, but there is penalty for GC<40%
    if 55 >= F >= 45 and 55 >= R >= 45:
        return 2.0
    elif 60 >= F >= 40 and 60 >= R >= 40:
        return 1.0
    elif F < 40 and R < 40:
        return -1.5 # it is very bad if both are less than 40%
    # penalize for GC<40%
    elif F < 40 or R < 40: # it is not that bad if only one is less than 40%
        if 60 >= F >= 40 or 60 >= R >= 40: 
            # and if the other one is in range it is not that bad
            return 0
        else:
            return -0.5
    else:
        return 0


# for mips
# F is the ligation arm
# R is the extension arm

# Scoring based on data from "Li JB, Genome Res, 2009" and "Turner EH, Nature Met, 2009"
def mip_TmScore(F, R): # max score 2
    # ref Li Genome Res 2009
    # for F range is 53-61 optimum is 58 (ligtaion arm)
    # for R range is 50-58 optimum is 53 (extension arm)
    if F > R:
        if 51<=R<=55 and 56<=F<=60:
            return 2.0
        elif 50<=R<=58 and 53<=F<=61:
            return 1.0
        else:
            return 0.5
    else:
        return 0


def mip_GCScore(F, R): # max score 2
    # Teer Genome Res 2010
    Score= 0.0
    if 45 <= R <= 55: 
        Score+= 1.0
    elif 40 <= R <= 60: 
        Score+= 0.5
    if 35 <= F <= 50: 
        Score+= 1.0
    elif 40 <= F <= 60: 
        Score+= 0.5
    return Score


def pcr_score_pair_IR(Line, path, Smax, even):
    # score for Tm max=2
    score= pcr_TmScore(Line[3], Line[13])
    # score for GC max=2
    score+= pcr_GCScore(Line[4], Line[14])
    score+= Penalty(Line[2], Line[2])
    # score for amplicon size, when everything else is the same
    # favor larger amplicons (max= 0.5)
    score+= 5**(1.0 * Line[21] / Smax) / 10
    score+= LenScore(Line[5], Line[15])
    return score


def pcr_score_pair_S(Line, path, Smax, even):
    # score for repeats, max 2 (1 for each)
    score= check_repeats(Line[2], Line[12], path)
    # score for Tm max=2
    score= pcr_TmScore(Line[3], Line[13])
    # score for GC max=2
    score+= pcr_GCScore(Line[4], Line[14])
    score+= Penalty(Line[2], Line[2])
    # score for amplicon size, when everything else is the same
    # favor larger amplicons (max= 0.5)
    score+= 5**(1.0 * Line[21] / Smax) / 10
    score+= LenScore(Line[5], Line[15])
    return score

       
def mip_picker(s, REdict, PriF, Middle, PriR, Arm_max, pntcutoff, pick_reverse, TTemp2, even, path, Picked_code, gfserver, Remove=False):
    
    if even:
        strand= "plus"
    else:
        strand= "minus"

    # score the picked list
    # here implement dimer check for the pair
    # each F and R primer is dimer checked during design which is the fastest
    # since number of primer is reduced upfront
    #s= [[mip_score_pair(p, even)] + p[1:] for p in s if p[0] != "na"]
    #s= [[mip_score_pair(p, even)] + p[1:] for p in s]

    # correct for repeat_score
    
    s.sort(reverse= True)
    
    for pair in s:
        goodpair= True
        if pair[11] not in Picked_code and pair[1] not in Picked_code:
            MIP= Filler(pair[2], pair[12], even, REdict, PriF, Middle, PriR, Arm_max)                  
            if MIP != "FAIL":
                if dimer(MIP, MIP, pntcutoff=2.5) == 'PASS':
                    Picked_code[pair[1]]= 0
                    Picked_code[pair[11]]= 0
                    thepair= pair[:] + [MIP, strand]

                    if not gfserver == "false":
                        # in-silico PCR
                        # this does not work on my mac only on linux
                        # you can define -maxSize=max_len * 5 (default=4000)
                        # which should be ok
                        proc = subprocess.Popen("%s %s %s stdout -out=bed" %(gfserver, pair[2], pair[12]), shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                        if not len([p.strip().split() for p in proc.stdout]) == 1:
                            # there has to be one and only one match
                            goodpair = False
                    elif Remove:
                        goodpair= False
                        repeat_score= check_repeats(pair[2], pair[12], path)
                        if repeat_score > 0:
                            # if it is 0 that means a very bad repeat score so discard 
                            pair[0]+= repeat_score
                            goodpair= True
                    if goodpair:
                        if pick_reverse == None:
                            keep= True
                            for olds in TTemp2[:-1]:
                                if olds[9] < thepair[8] < thepair[9] < olds[18] < olds[19] < thepair[18] or olds[19] < thepair[8] or thepair[9] < olds[8] < olds[19] < thepair[18]:
                                    keep= True
                                else:
                                    keep= False
                                    break
                            if keep:
                                return Picked_code, thepair
                        else:
                            return Picked_code, thepair
    return Picked_code, 'FAIL'


def score_mip_IR(Line, path, Smax, even):
    score= 5**(1.0 * Line[21] / Smax) / 10
    score+= Penalty(Line[2], Line[12])
    score+= LenScore(Line[5], Line[15])
    if even:
        # score for Tm max=2
        score+= mip_TmScore(Line[3], Line[13])
        # score for GC max=2
        score+= mip_GCScore(Line[4], Line[14])
    else:
        ## score for Tm max=2
        score+= mip_TmScore(Line[13], Line[3])
        # score for GC max=2
        score+= mip_GCScore(Line[14], Line[4])
    return score


def score_mip_S(Line, path, Smax, even):
    score= 5**(1.0 * Line[21] / Smax) / 10
    # score repeats
    score+= check_repeats(Line[2], Line[12], path)
    score+= Penalty(Line[2], Line[12])
    score+= LenScore(Line[5], Line[15])
    if even:
        # score for Tm max=2
        score+= mip_TmScore(Line[3], Line[13])
        # score for GC max=2
        score+= mip_GCScore(Line[4], Line[14])
    else:
        ## score for Tm max=2
        score+= mip_TmScore(Line[13], Line[3])
        # score for GC max=2
        score+= mip_GCScore(Line[14], Line[4])
    return score


def permutations(items, s):
    s= int(s)
    if s==0:
        yield ''
    else:
        for i in xrange(len(items)):  
            for item in permutations(items, s-1):  
                 yield str(items[i])+str(item)


def REcheck(s, REdict):
    s= s.lower()
    for keys in REdict:
        if s.count(REdict[keys][0]) + s.count(revcomp(REdict[keys][0])) != REdict[keys][1]:
            # fail
            return 0
    # pass
    return 1


def Filler(PrimerF, PrimerR, even, REdict, PriF, Middle, PriR, Arm_max):
    # make F+R=Arm_max and the whole MIP 100mer
    if even: #even
        Arm5= revcomp(PrimerF)
        Arm3= PrimerR
    else: #odd
        Arm5= revcomp(PrimerR)
        Arm3= PrimerF

    Arm5= Arm5.upper()
    Arm3= Arm3.upper()
    Larm3= len(Arm3)
    Larm5= len(Arm5)
    if Larm5 + Larm3 == Arm_max:
        pMIP= ''.join([PriF, Arm5, Middle, Arm3, PriR])
        if REcheck(pMIP, REdict):
            return pMIP
        else:
            return 'FAIL'

    else:
        # add bases at the end of Arm5 and begining of Arm3
        base_req= Arm_max-Larm5-Larm3
        arm_diff= abs(Larm5-Larm3)
        if Larm5 == Larm3:
            addto3= [p for p in permutations(bases, base_req/2)]
            addto5= addto3[:]
            shuffle(addto3)
            shuffle(addto5)
            for i in xrange(0, len(addto5)):
                for I in xrange(0, len(addto3)):
                    pMIP= ''.join([PriF, Arm5, addto5[i], Middle, addto3[I], Arm3, PriR])
                    if REcheck(pMIP, REdict):
                        return pMIP
            return 'FAIL'

        elif base_req <= arm_diff: # all bases go to the shorter arm
            if len(Arm5) > len(Arm3):
                # work on Arm3
                addto3= [p for p in permutations(bases, base_req)]
                shuffle(addto3)
                for I in xrange(0, len(addto3)):
                    pMIP= ''.join([PriF, Arm5, Middle, addto3[I], Arm3, PriR])
                    if REcheck(pMIP, REdict):
                        return pMIP
                return 'FAIL'

            elif len(Arm3) > len(Arm5):
                # work on Arm5
                addto5= [p for p in permutations(bases, base_req)]
                shuffle(addto5)
                for i in xrange(0, len(addto5)):
                    pMIP= ''.join([PriF, Arm5, addto5[i], Middle, Arm3, PriR])
                    if REcheck(pMIP, REdict):
                        return pMIP
                return 'FAIL'

        elif base_req > arm_diff: # distribute bases as necessary
            if len(Arm5) > len(Arm3):
                # work on Arm3, it is the shorter
                for3= (base_req+arm_diff)/2
                for5= base_req-for3
                addto3= [p for p in permutations(bases, for3)]
                addto5= [p for p in permutations(bases, for5)]
                shuffle(addto3)
                shuffle(addto5)
                for I in xrange(0, len(addto3)):
                    for i in xrange(0, len(addto5)):
                        pMIP= ''.join([PriF, Arm5, addto5[i], Middle, addto3[I], Arm3, PriR])
                        if REcheck(pMIP, REdict):
                            return pMIP
                return 'FAIL'

            elif len(Arm3) > len(Arm5):
                # work on Arm5, it is the shorter
                for5= (base_req+arm_diff)/2
                for3= base_req-for5
                addto3= [p for p in permutations(bases, for3)]
                addto5= [p for p in permutations(bases, for5)]
                shuffle(addto3)
                shuffle(addto5)
                for I in xrange(0, len(addto5)):
                    for i in xrange(0, len(addto3)):
                        pMIP= ''.join([PriF, Arm5, addto5[I], Middle, addto3[i], Arm3, PriR])
                        if REcheck(pMIP, REdict):
                            return pMIP
                return 'FAIL'
